import datetime
import json
import os
import traceback
from urllib.parse import urlparse, parse_qs

import elasticsearch
import geoip2.database
import numpy as np
import pandas as pd
from elasticsearch import Elasticsearch
from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn

from base.util.common_util import get_uri_depth, get_uri_filename_length, get_project_root, build_expression
from base.constant import out_folder
from base.util.file_util import check_path
from base.util.geo_util import get_geo_by_ip
from base.util.log_util import log_e


def get_url_param_count(url):
    query = urlparse(url).query  # 解析 URL 中的查询字符串
    params = parse_qs(query)  # 解析查询字符串为字典
    return len(params)


def save_session_id(origin_list, session_id_file):
    res = []
    reader = geoip2.database.Reader(f'{get_project_root()}base/util/GeoLite2-City.mmdb')
    for item in origin_list:
        _source = item.get("_source", {})
        source = _source.get("source", {})
        tcpflags = _source.get("tcpflags", {})
        destination = _source.get("destination", {})
        http = _source.get("http", {})
        dns = _source.get("dns", {})
        tls = _source.get("tls", {})
        uri = http.get('uri', [])
        uri_length = [len(u) for u in uri]
        uri_depth = [get_uri_depth(u) for u in uri]
        uri_filename_length = [get_uri_filename_length(u) for u in uri]
        uri_params = [get_url_param_count(u) for u in uri]
        newItem = {
            "id": item["_id"],
            "node": _source.get("node", ""),
            "segmentCnt": _source.get("segmentCnt", 0),
            "tcpflags.rst": tcpflags.get("rst", 0),
            "tcpflags.ack": tcpflags.get("ack", 0),
            "tcpflags.syn": tcpflags.get("syn", 0),
            "tcpflags.urg": tcpflags.get("urg", 0),
            "tcpflags.psh": tcpflags.get("psh", 0),
            "tcpflags.syn-ack": tcpflags.get("syn-ack", 0),
            "tcpflags.fin": tcpflags.get("fin", 0),
            "source.ip": source.get("ip", ""),
            "destination.ip": destination.get("ip", ""),
            "source.port": source.get("port", ""),
            "source.packets": source.get("packets", ""),
            "source.bytes": source.get("bytes", 0),
            "destination.port": destination.get("port", ""),
            "destination.bytes": destination.get("bytes", 0),
            "destination.packets": destination.get("packets", 0),
            "initRTT": _source.get("initRTT", ""),
            "firstPacket": _source.get("firstPacket", 0),
            "lastPacket": _source.get("lastPacket", 0),
            "ipProtocol": _source.get("ipProtocol", ""),
            "protocolCnt": _source.get("protocolCnt", 0),
            "protocol": _source.get("protocol", ""),
            "server.bytes": _source.get("server", {}).get("bytes", 0),
            "totDataBytes": _source.get("totDataBytes", 0),
            "network.packets": _source.get("network", {}).get("packets", 0),
            "network.bytes": _source.get("network", {}).get("bytes", 0),
            "length": _source.get("length", 0),
            "client.bytes": _source.get("client", {}).get("bytes", 0),
            "http.uri": uri,
            "http.uri_length_mean": round(np.nan_to_num(np.mean(uri_length)), 5),
            "http.uri_length_var": round(np.nan_to_num(np.var(uri_length)), 5),
            "http.uri_param_count_mean": round(np.nan_to_num(np.mean(uri_params)), 5),
            "http.uri_param_count_var": round(np.nan_to_num(np.var(uri_params)), 5),
            "http.uri_depth_mean": round(np.nan_to_num(np.mean(uri_depth)), 5),
            "http.uri_depth_var": round(np.nan_to_num(np.var(uri_depth)), 5),
            "http.uri_filename_length_mean": round(np.nan_to_num(np.mean(uri_filename_length)), 5),
            "http.uri_filename_length_var": round(np.nan_to_num(np.var(uri_filename_length)), 5),

            "http.response-content-type": http.get("response-content-type", []),
            "http.bodyMagicCnt": http.get("bodyMagicCnt", 0),
            "http.statuscodeCnt": http.get("statusCodeCnt", 0),
            "http.clientVersionCnt": http.get("clientVersionCnt", 0),
            "http.response-content-typeCnt": http.get("response-content-typeCnt", 0),
            "http.xffIpCnt": http.get("xffIpCnt", 0),
            "http.requestHeaderCnt": http.get("requestHeaderCnt", 0),
            "http.serverVersion": http.get("serverVersion", []),
            "http.responseHeaderCnt": http.get("responseHeaderCnt", 0),
            "http.xffIp": http.get("xffIp", []),
            "http.clientVersion": http.get("clientVersion", []),
            "http.uriTokens": http.get("uriTokens", ""),
            "http.useragentCnt": http.get("useragentCnt", 0),
            "http.statuscode": http.get("statusCode", []),
            "http.bodyMagic": http.get("bodyMagic", []),
            "http.request-content-type": http.get("request-content-type", []),
            "http.uriCnt": http.get("uriCnt", 0),
            "http.serverVersionCnt": http.get("serverVersion", []),
            "http.useragent": http.get("useragent", []),
            "http.keyCnt": http.get("keyCnt", 0),
            "http.request-referer": http.get("requestReferer", []),
            "http.request-refererCnt": http.get("requestRefererCnt", 0),
            "http.path": http.get("path", []),
            "http.hostCnt": http.get("hostCnt", 0),
            "http.response-server": http.get("response-server", []),
            "http.pathCnt": http.get("pathCnt", 0),
            "http.useragentTokens": http.get("useragentTokens", ""),
            "http.methodCnt": http.get("methodCnt", 0),
            "http.method": http.get("method", []),
            "http.method-GET": http.get("method-GET", 0),
            "http.method-POST": http.get("method-POST", 0),
            "http.key": http.get("key", []),
            "http.hostTokens": http.get("hostTokens", ""),
            "http.requestHeader": http.get("requestHeader", []),
            "http.responseHeader": http.get("responseHeader", []),
            "dns.ASN": dns.get("ASN", []),
            "dns.RIR": [item for item in dns.get("RIR", []) if item != ""],
            "dns.GEO": dns.get("GEO", []),
            "dns.alpn": dns.get("https.alpn", []),
            "dns.alpnCnt": dns.get("https.alpnCnt", 0),
            "dns.ip": dns.get("ip", []),
            "dns.ipCnt": dns.get("ipCnt", 0),
            "dns.OpCode": dns.get("opcode", []),
            "dns.OpCodeCnt": dns.get("opcodeCnt", 0),
            "dns.Puny": dns.get("puny", []),
            "dns.PunyCnt": dns.get("puntCnt", 0),
            "dns.QueryClass": dns.get("qc", []),
            "dns.QueryClassCnt": dns.get("qcCnt", 0),
            "dns.QueryType": dns.get("qt", []),
            "dns.QueryTypeCnt": dns.get("qtCnt", 0),
            "dns.status": dns.get("status", []),
            "dns.hostCnt": json.dumps(dns.get("hostCnt", 0)),
            "dns.host": json.dumps(dns.get("host", [])),
            "dns.statusCnt": dns.get("statusCnt", 0),
            "tls.cipher": tls.get("cipher", []),
            "tls.cipherCnt": tls.get("cipherCnt", 0),
            "tls.dstSessionId": tls.get("dstSessionId", []),
            "tls.ja3": tls.get("ja3", []),
            "tls.ja3Cnt": tls.get("ja3Cnt", 0),
            "tls.ja3s": tls.get("ja3s", []),
            "tls.ja3sCnt": tls.get("ja3sCnt", 0),
            "tls.ja4": tls.get("ja4", []),
            "tls.ja4Cnt": tls.get("ja4Cnt", 0),
            "tls.srcSessionId": tls.get("srcSessionId", []),
            "tls.version": tls.get("version", []),
            "tls.versionCnt": tls.get("versionCnt", 0),
            "tls.ja4_r": tls.get("versionCnt", 0),
            "tls.ja4_rCnt": tls.get("versionCnt", 0),
            "packetPos": json.dumps(_source.get("packetPos", []))
        }
        res.append(get_geo_by_ip(reader, newItem))
        session_id_file = check_path(session_id_file)
    file_exists = os.path.isfile(session_id_file)
    write_header = not file_exists or os.path.getsize(session_id_file) == 0
    pd.DataFrame(res).to_csv(check_path(session_id_file), mode='a', header=write_header, index=False)
    # if len(res) == 0:
    #     print(f"[session]写入0条，说明请求有缺失，请删除本地缓存后重新执行")
    #     exit(0)


class LogTextColumn(TextColumn):
    def render(self, task):
        return task.fields.get("custom_text", "--|--")


def get_session(config_yaml):
    http_auth = (config_yaml['es']['account'], config_yaml['es']['password']) if config_yaml['es']['use_http_auth'] \
        else None
    es = Elasticsearch(config_yaml['es']['url'],
                       http_auth=http_auth
                       , verify_certs=False)
    print("[session]清理滚动id...")
    es.clear_scroll(body={'scroll_id': '_all'})
    start_session_time = datetime.datetime.now()
    progress = Progress(TextColumn("[progress.description]{task.description}"),
                        BarColumn(),
                        TaskProgressColumn(),
                        LogTextColumn("[progress.description]{task.description}"), )
    progress.start()
    session_id_file = f"{out_folder}{config_yaml['session']['time_session_out']}"
    try:
        print("[session]开始搜索")
        arkime_expression = config_yaml['session']['expression']
        count_exp = build_expression(size=None,
                                     start_time=config_yaml['session']["start_time"],
                                     end_time=config_yaml['session']["end_time"],
                                     arkime_expression=arkime_expression)
        print("[session]开始根据时间范围获取数量")
        total = es.count(body=count_exp)['count']
        print(f"[session]获取数量:{total}")
        task_id = progress.add_task("获取session", total=total)
        scroll_exp = build_expression(size=config_yaml['session']['size'],
                                      start_time=config_yaml['session']["start_time"],
                                      end_time=config_yaml['session']["end_time"],
                                      arkime_expression=arkime_expression)
        first_result = es.search(body=scroll_exp, scroll="1m")
        scroll_id = first_result.body['_scroll_id']
        first_list = first_result.body['hits']['hits']
        if total == 0:
            print("[session]未获取到session")
            exit(0)
        progress.update(task_id, completed=len(first_list), custom_text=f"当前获取：{len(first_list)}/{total}")
        progress.refresh()
        save_session_id(origin_list=first_list, session_id_file=session_id_file)
        resul_counts = len(first_list)
        print(f"[session]当前获取{resul_counts}")
        while True:
            now = datetime.datetime.now()
            try:
                result = es.scroll(scroll_id=scroll_id, scroll=str(config_yaml['session']["session_alive"]),
                                   request_timeout=120)
            except Exception as e:
                print(f"[session]查询出错：{e},将继续查询")
                continue
            hits = result.body['hits']['hits']
            save_session_id(origin_list=hits, session_id_file=session_id_file)
            if len(hits) != 0:
                resul_counts += len(hits)
                progress.update(task_id, completed=resul_counts,
                                custom_text=f"当前获取{resul_counts}/{total}条,耗时：{str(datetime.datetime.now() - now)},总耗时：{str(datetime.datetime.now() - start_session_time)}")
                progress.refresh()
            else:
                break
        print(f"[session]一共获取{resul_counts}条，保存在{session_id_file}")
        es.clear_scroll(body={"scroll_id": scroll_id})
        progress.stop()
    except elasticsearch.NotFoundError:
        log_e("[session]分片没找到")
        progress.stop()
        exit(0)
    except Exception as e:
        print(f"[session]错误：{e}")
        traceback.print_exc()
        progress.stop()
        exit(0)
    print("[session]获取session完毕，记得启动 arkime capture")
